# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: xguo

import time
import collections
import logging
import math
from concurrent.futures import ThreadPoolExecutor

import tornado.ioloop
import tornado.testing
from tornado import gen
from tornado.ioloop import PeriodicCallback
from tornado.testing import AsyncTestCase
from tornado.concurrent import run_on_executor

import coin.proto.coin_order_enums_pb2 as coin_order
from coin.exchange.base.order_gateway import OrderSide, OrderType
from coin.flow.simple_queue import SimpleQueue
from coin.flow.subscriber import from_queue
from coin.proto.coin_query_pb2 import ProductOrderElement
from coin.strategy.hk_hard_arb.og_util import print_working_orders
from coin.exchange.okex_futures.types import OkexFuturesOrderSide

OGFixture = collections.namedtuple('OGFixture',
                                   [
                                       'order_gateway',
                                       'market_type',
                                       'currencies',
                                       'product',
                                       'public_client',
                                       'private_client',
                                       'config',
                                       'gen_order',
                                       'feed_subsystem_cls',
                                   ])


def to_proto_order_side(order_side):
  if isinstance(order_side, OrderSide):
    order_side_map = {
        OrderSide.BUY: coin_order.BUY_ORDER,
        OrderSide.SELL: coin_order.SELL_ORDER,
    }
  elif isinstance(order_side, OkexFuturesOrderSide):
    order_side_map = {
        OkexFuturesOrderSide.BUY_OPEN: coin_order.BUY_OPEN_ORDER,
        OkexFuturesOrderSide.SELL_OPEN: coin_order.SELL_OPEN_ORDER,
    }
  else:
    raise ValueError('Invalid input type: %s' % type(order_side))
  res = order_side_map[order_side]
  return res


def from_proto_order_side(order_side: int):
  order_side_map = {
      coin_order.BUY_ORDER: OrderSide.BUY,
      coin_order.SELL_ORDER: OrderSide.SELL,
      coin_order.BUY_OPEN_ORDER: OkexFuturesOrderSide.BUY_OPEN,
      coin_order.SELL_OPEN_ORDER: OkexFuturesOrderSide.SELL_OPEN,
  }
  if order_side in order_side_map:
    return order_side_map[order_side]
  else:
    raise ValueError(order_side)


def is_order_in_list(orig_order, order_list):
  for order in order_list:
    price_same = math.isclose(order.price, orig_order.price)
    qty_same = math.isclose(order.qty, orig_order.qty)
    side_same = (order.side == orig_order.side)
    if price_same and qty_same and side_same:
      return True
  return False


def same_order_lists(order_list1, order_list2):
  if len(order_list1) != len(order_list2):
    return False

  for order in order_list1:
    if not is_order_in_list(order, order_list2):
      return False

  return True


def working_orders_to_order_list(working_orders):
  order_list = []
  for order in working_orders:
    order_proto = ProductOrderElement(price=order.price,
                                      qty=order.qty,
                                      side=to_proto_order_side(order.order_side))
    order_list.append(order_proto)
  return order_list


def find_balance(currency, account_balance):
  for balance in account_balance.each_balance:
    if balance.currency == currency.currency:
      return balance.available


class BookReceiver(object):
  def __init__(self, fixture, book_callback):
    self.fixture = fixture
    self.feed_subsystem = fixture.feed_subsystem_cls(fixture.product, self.on_book_reset)
    self.ioloop = None
    self.book_callback = book_callback

  def on_book_reset(self, exchange, flow_book_builder):
    flow_book_builder.subscribe_book(self.fixture.product, self.on_book)

  def start(self, ioloop):
    queue = SimpleQueue('hk_og_test')
    flow_sub = from_queue(queue)
    self.ioloop = ioloop
    self.feed_subsystem.start(ioloop, queue, flow_sub)

  def on_book(self, book):
    self.book_callback(book)


class PassAggOrderTester(object):
  def __init__(self, fixture, mode, max_num_orders):
    assert mode in ('passive', 'aggressive')
    self.fixture = fixture
    self.feed_subsystem = fixture.feed_subsystem_cls(fixture.product, self.on_book_reset)
    self.mode = mode
    self.num_orders = 0
    self.max_num_orders = max_num_orders
    self.ioloop = None
    self.last_book_time = 0

  def on_book_reset(self, exchange, flow_book_builder):
    flow_book_builder.subscribe_book(self.fixture.product, self.on_book)

  def start(self, ioloop):
    queue = SimpleQueue('hk_og_test')
    flow_sub = from_queue(queue)
    self.ioloop = ioloop
    self.feed_subsystem.start(ioloop, queue, flow_sub)

  def on_book(self, book):
    if time.time() - self.last_book_time < 2:
      return

    if self.num_orders > self.max_num_orders:
      return

    og = self.fixture.order_gateway
    mid_price = 0.5 * (book.ask0().price + book.bid0().price)
    orders = self.fixture.gen_order(mid_price, self.mode)
    for order in orders:
      og.submit(self.fixture.product,
                order.price,
                order.qty,
                from_proto_order_side(order.side),
                OrderType.LIMIT)
    self.num_orders += len(orders)
    self.last_book_time = time.time()
    self.ioloop.call_later(1, og.cancel_all)


def gen_og_test_class(fixture_func):
  class TestOG(AsyncTestCase):
    def setUp(self):
      super().setUp()
      queue = SimpleQueue('queue')
      self.fill_messages = []
      self.fixture = fixture_func()
      self.fixture.order_gateway._fill_callbacks.append(self.on_fill)
      self.fixture.order_gateway.start(queue, self.io_loop)
      # Make timeout bigger since dev machine's connection is often slow.
      self.fixture.public_client.npubc.timeout = 20
      self.fixture.private_client.nprivc.timeout = 20
      self.latest_ticker = None
      self.book = None
      self.open_orders = None

      self.periodic_updater = PeriodicCallback(self.update_info, 3000)
      self.executor = ThreadPoolExecutor(max_workers=1)
      self.logger = logging.getLogger('OG_test')
      self.io_loop.set_blocking_log_threshold(0.1)

    def tearDown(self):
      super().tearDown()
      if self.periodic_updater.is_running():
        self.periodic_updater.stop()
      self.fixture = None

    def on_book(self, book):
      self.book = book

    @run_on_executor
    def update_info(self):
      product = self.fixture.product
      try:
        update = self.fixture.private_client.query_list_orders(product)
        self.open_orders = update.msg
      except Exception:
        self.logger.exception('query open order failed!')
      print_working_orders(self.logger, self.fixture.order_gateway)

    def on_fill(self, price, qty, order):
      self.fill_messages.append((price, qty, order))

    def ensure_working_orders_accepted(self):
      all_accepted = False
      for _ in range(20):
        all_accepted = all(
            order.accepted for order in self.fixture.order_gateway.get_working_order())
        yield gen.sleep(1)
      self.assertTrue(all_accepted)

    def ensure_same_order_list(self, orders):
      og = self.fixture.order_gateway
      working_order_list = working_orders_to_order_list(og.get_working_order())
      if not same_order_lists(working_order_list, orders):
        print_working_orders(self.logger, self.fixture.order_gateway)
        self.logger.info(self.open_orders)
        self.logger.info(working_order_list)
        self.logger.info(orders)

      self.assertTrue(same_order_lists(working_order_list, orders))

      # Ensure that orders are indeed accepted by the exchange.
      for _ in range(20):
        if not same_order_lists(working_order_list, self.open_orders.each_order):
          print_working_orders(self.logger, self.fixture.order_gateway)
          yield gen.sleep(1)
      if not same_order_lists(working_order_list, self.open_orders.each_order):
        self.logger.info('The OG working_order_list and rest query order list are different!')
        print_working_orders(self.logger, self.fixture.order_gateway)
        self.logger.info(self.open_orders)
      self.assertTrue(same_order_lists(working_order_list, self.open_orders.each_order))

    def wait_until_og_ready(self):
      while not self.fixture.order_gateway.is_ready():
        yield gen.sleep(1)

      while len(self.fixture.order_gateway.get_working_order()) > 0:
        self.fixture.order_gateway.cancel_all()
        yield gen.sleep(1)

      if self.periodic_updater.is_running():
        while self.book is None or self.book.ask0() is None or self.book.bid0(
        ) is None or self.open_orders is None:
          yield gen.sleep(1)

    def cleanup_open_orders(self):
      # Cancel open orders.
      og = self.fixture.order_gateway
      for n in range(20):
        self.logger.info('Cancel working orders: %s', n)
        og.cancel_all()
        yield gen.sleep(5)  # Wait for a while until cancel is successful.
        if len(og.get_working_order()) == 0:
          break

      self.assertEqual(len(og.get_working_order()), 0)

      for n in range(20):
        if len(self.open_orders.each_order) == 0:
          break
        else:
          self.logger.info('working orders by rest query is not cleared. %s', n)
        yield gen.sleep(5)  # Wait for a while until cancel is successful.
      self.assertEqual(len(self.open_orders.each_order), 0)

    def gen_big_order(self):
      og = self.fixture.order_gateway
      product = self.fixture.product

      account_balance = og._og_info.balance_as_proto()
      available = find_balance(product.base, account_balance)

      # price use passive order's sell price, qty is a value that is bigger than
      # available balance, and also satisfies precision requirement.
      orders = self.fixture.gen_order(self.latest_ticker.last, 'passive')
      sell_order = None
      for order in orders:
        if order.side == coin_order.SELL_ORDER:
          sell_order = order
          break

      self.assertTrue(sell_order is not None)
      qty = sell_order.qty
      while qty < available * 1.1:
        qty += qty
      return sell_order.price, qty

    @tornado.testing.gen_test(timeout=20)
    def test_get_balances(self):
      og = self.fixture.order_gateway
      yield from self.wait_until_og_ready()

      account_balance = og.og_info.balance_as_proto()
      self.assertTrue(len(list(account_balance.each_balance)) > 0)
      currency_str_list = [currency.currency for currency in self.fixture.currencies]
      for balance in account_balance.each_balance:
        if balance.currency in currency_str_list:
          print(balance)

    @tornado.testing.gen_test(timeout=60)
    def test_cancel_fail(self):
      og = self.fixture.order_gateway
      product = self.fixture.product
      yield from self.wait_until_og_ready()

    @tornado.testing.gen_test(timeout=120)
    def test_submit_fail(self):
      self.periodic_updater.start()
      og = self.fixture.order_gateway
      product = self.fixture.product
      yield from self.wait_until_og_ready()

      price, qty = self.gen_big_order()
      # Insufficient balance error is expected.
      og.submit(product, price, qty, OrderSide.SELL, OrderType.LIMIT)
      yield gen.sleep(5)

      open_orders = list(self.open_orders.each_order)

      # The submission is failed, and there is no open_order.
      self.assertEqual(len(open_orders), 0)

      # Failed order will be expunged from working orders, after receiving a
      # submission error message, or after 60 seconds.
      wait_sec = 0
      while og.get_working_order():
        working_orders = list(og.get_working_order())
        self.assertEqual(len(working_orders), 1)
        self.assertFalse(working_orders[0].accepted)
        yield gen.sleep(1)
        wait_sec += 1
        if wait_sec == 70:
          break
      self.assertLess(wait_sec, 70, 'unaccepted order not purged after 70 seconds.')

    @tornado.testing.gen_test(timeout=400)
    def test_pass_order(self):
      book_receiver = BookReceiver(self.fixture, self.on_book)
      book_receiver.start(self.io_loop)
      self.periodic_updater.start()
      og = self.fixture.order_gateway
      product = self.fixture.product
      yield from self.wait_until_og_ready()

      # OG is ready.
      last = (self.book.ask0().price + self.book.bid0().price) / 2
      orders = self.fixture.gen_order(last, 'passive')
      for order in orders:  # buy and sell order.
        og.submit(product,
                  order.price,
                  order.qty,
                  from_proto_order_side(order.side),
                  OrderType.LIMIT)
        yield gen.sleep(2)

      # Verify OG's working_orders
      yield from self.ensure_working_orders_accepted()
      # Users == OG.working_orders == exchange.working_orders
      yield from self.ensure_same_order_list(orders)
      # Cancel all open orders.
      yield from self.cleanup_open_orders()

    @tornado.testing.gen_test(timeout=1200)
    def test_intense_pass_order(self):
      book_receiver = BookReceiver(self.fixture, self.on_book)
      book_receiver.start(self.io_loop)
      self.periodic_updater.start()
      og = self.fixture.order_gateway
      product = self.fixture.product
      yield from self.wait_until_og_ready()

      for _ in range(100):
        # OG is ready.
        last = (self.book.ask0().price + self.book.bid0().price) / 2
        orders = self.fixture.gen_order(last, 'passive')
        for order in orders:  # buy and sell order.
          og.submit(product,
                    order.price,
                    order.qty,
                    from_proto_order_side(order.side),
                    OrderType.LIMIT)
          yield gen.sleep(0.5)
        yield gen.sleep(1)
        og.cancel_all()
        yield gen.sleep(1)

      # Cancel all open orders.
      yield from self.cleanup_open_orders()

    @tornado.testing.gen_test(timeout=60)
    def test_agg_order(self):
      book_receiver = BookReceiver(self.fixture, self.on_book)
      book_receiver.start(self.io_loop)
      self.periodic_updater.start()
      og = self.fixture.order_gateway
      product = self.fixture.product
      yield from self.wait_until_og_ready()

      for _ in range(5):  # Try 5 times.
        last = (self.book.ask0().price + self.book.bid0().price) / 2
        orders = self.fixture.gen_order(last, 'aggressive')
        for order in orders:
          og.submit(product,
                    order.price,
                    order.qty,
                    from_proto_order_side(order.side),
                    OrderType.LIMIT)
          yield gen.sleep(1)

        open_orders = list(self.open_orders.each_order)
        # Ensure that if there is any listed open orders, it must be orders
        # that we just placed.
        for order in open_orders:
          self.assertTrue(is_order_in_list(order, orders))

        # Check whether fill message is received.
        if len(self.fill_messages) > 0:
          break
        yield gen.sleep(5)

      self.assertGreater(len(self.fill_messages), 0)
      # Cancel all open orders.
      yield from self.cleanup_open_orders()

    @tornado.testing.gen_test(timeout=120)
    def test_unacked_order(self):
      og = self.fixture.order_gateway
      product = self.fixture.product
      public_client = self.fixture.public_client
      yield from self.wait_until_og_ready()

      # OG is ready.
      update = public_client.query_ticker(product)
      orders = self.fixture.gen_order(update.msg.last, 'passive')
      for order in orders:
        og.submit(product,
                  order.price,
                  order.qty,
                  from_proto_order_side(order.side),
                  OrderType.LIMIT)
      yield gen.sleep(5)
      for order in og.get_working_order():
        self.assertTrue(order.accepted)
        order.accepted = False  # Forcefully set it to be False

      # Many long stay unacked order should result in OG restarting.
      yield gen.sleep(60)

      # OG should finish restarting and restore all the open orders.
      self.assertEqual(len(og.get_working_order()), len(orders))
      self.ensure_working_orders_accepted()
      self.cleanup_open_orders()

  return TestOG
